import numpy as np
import pandas as pd
import os
import argparse

def process_channel_data(input_csv, output_dir, window_size=20, overlap=10, use_all_clusters=True):
    """
    处理信道数据，将连续时间点的数据分割成有重叠的窗口
    
    参数:
        input_csv: 输入CSV文件路径
        output_dir: 输出目录
        window_size: 窗口大小(每个输入的时间点数量)
        overlap: 相邻窗口的重叠时间点数量
        use_all_clusters: 是否使用所有簇
    """
    print(f"读取数据: {input_csv}")
    df = pd.read_csv(input_csv)
    
    # 查看输入数据的基本信息
    total_time_points = df['receiver_point'].nunique()
    total_paths_per_point = df['path_number'].nunique()
    print(f"总时间点数: {total_time_points}")
    print(f"每个时间点的路径数: {total_paths_per_point}")
    
    # 使用所有可用的簇
    n_clusters = total_paths_per_point
    print(f"使用所有{n_clusters}个簇")
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 将数据组织成3D数组: [time_points, clusters, features]
    time_points = sorted(df['receiver_point'].unique())
    
    # 创建空数组
    all_features = np.zeros((len(time_points), n_clusters, 5))
    
    # 填充数据
    print("组织数据...")
    for t_idx, t in enumerate(time_points):
        # 获取当前时间点的数据
        time_data = df[df['receiver_point'] == t]
        # 按路径编号排序，确保顺序一致
        sorted_paths = time_data.sort_values('path_number')
        
        # 检查是否有足够的路径
        if len(sorted_paths) < n_clusters:
            print(f"警告: 时间点{t}只有{len(sorted_paths)}个路径，少于期望的{n_clusters}个")
        
        # 填充特征数组
        for c_idx, (_, path) in enumerate(sorted_paths.iterrows()):
            if c_idx >= n_clusters:
                break
            # 按照模型要求的顺序: [horizontal_aoa, vertical_aoa, power, ASA, ZSA]
            all_features[t_idx, c_idx, 0] = path['horizontal_aoa']  # 水平到达角
            all_features[t_idx, c_idx, 1] = path['vertical_aoa']    # 垂直到达角
            all_features[t_idx, c_idx, 2] = path['received_power']  # 功率
            all_features[t_idx, c_idx, 3] = path['ASA']             # 水平扩展角
            all_features[t_idx, c_idx, 4] = path['ZSA']             # 垂直扩展角
    
    # 创建滑动窗口
    print("创建滑动窗口...")
    windows = []
    window_indices = []
    step = window_size - overlap
    
    for i in range(0, len(time_points) - window_size + 1, step):
        window = all_features[i:i+window_size]
        windows.append(window)
        window_indices.append((i, i+window_size))
    
    print(f"共创建了{len(windows)}个窗口")
    
    # 保存为npz格式
    output_path = os.path.join(output_dir, 'channel_data.npz')
    print(f"保存数据到: {output_path}")
    
    # 转换为最终格式
    # Windows可能有不同长度，将所有窗口拼接成一个大数组
    channel_features = np.vstack([w.reshape(1, window_size, n_clusters, 5) for w in windows])
    
    # 保存数据
    np.savez(
        output_path,
        channel_features=channel_features,  # [n_windows, window_size, n_clusters, 5]
        window_indices=window_indices       # 记录每个窗口的原始索引
    )
    
    print("数据处理完成!")
    print(f"特征形状: {channel_features.shape}")
    
    return output_path

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='处理信道数据CSV文件')
    parser.add_argument('--input', type=str, required=True, help='输入CSV文件的路径')
    parser.add_argument('--output_dir', type=str, default='./dataset/channel_data', help='输出目录')
    parser.add_argument('--window_size', type=int, default=20, help='窗口大小(每个输入的时间点数量)')
    parser.add_argument('--overlap', type=int, default=10, help='相邻窗口的重叠时间点数量')
    
    args = parser.parse_args()
    
    process_channel_data(
        args.input,
        args.output_dir,
        args.window_size,
        args.overlap,
        True
    ) 